#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/11/9 14:40
# @Author  : CHEN Wang
# @Site    : 
# @File    : factor_return_related.py
# @Software: PyCharm

"""
脚本说明: 获取因子收益相关数据api
"""

import copy

from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import BOND_4_FACTOR, BOND_7_FACTOR, \
    FAMA_3_FACTOR, FAMA_5_FACTOR, BARRA_STYLE_FACTOR, BARRA_INDUSTRY_FACTOR, BARRA_FULL_FACTOR
from quant_researcher.quant.project_tool.db_operator import my_mysql, db_conn
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.wrapper_tools import common_wrappers
from quant_researcher.quant.project_tool.logger.my_logger import LOG

T_BOND_FACTOR_RETURN = 'factors_bond_stylefactorsret'
C_DATE = 'trade_date'

T_FAMA_FACTOR_RETURN = 'fama_french_factorsret'
T_BARRA_FACTOR_RETURN = 'stk_barra_factor_ret'
EARLIEST_START_DATE = '1900-01-01'


def get_bond_factor_return(factor_type, start_date=None, end_date=None, only_latest=False):
    """
    获取债券因子收益率数据

    :param str factor_type: BOND_4YZ-债券四因子；BOND_7YZ-债券七因子
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame，前四个因子为债券四因子
                - end_date: 日期
                - Level: 系统因子
                - Slope: 期限因子
                - Convertible: 权益因子
                - Credit: 信用因子
                - Convex: 利率曲线凸度因子
                - Default: 违约因子
                - Currency: 货币因子
    """
    if factor_type == 'BOND_4YZ':
        select = copy.deepcopy(BOND_4_FACTOR)
    elif factor_type == 'BOND_7YZ':
        select = copy.deepcopy(BOND_7_FACTOR)
    else:
        raise NotImplementedError('目前只有4因子和7因子')
    select.insert(0, 'trade_date')

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    where = []

    conn = db_conn.get_tk_factors_conn()

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=T_BOND_FACTOR_RETURN, select='max(tradedate)',
                                       where=f"tradedate <= '{end_date}'",
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{T_BOND_FACTOR_RETURN}表为空，请检查")
            return
        where.append(f"trade_date = '{latest_date}'")
    else:
        where.append(f"trade_date >= '{start_date}'")
        where.append(f"trade_date <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=T_BOND_FACTOR_RETURN, select=select, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{factor_type}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df = df.rename(columns={'trade_date': 'end_date'})
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_fama_factor_return(factor_type, start_date=None, end_date=None, only_latest=False):
    """
    获得FAMA因子收益率数据

    :param str factor_type: FAMA_3YZ-3因子；FAMA_5YZ-5因子
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame，前三个为FAMA三因子
                - end_date: 日期
                - rm_freeshares: 风险溢价因子
                - smb_freeshares: 规模因子
                - hml_freeshares: 估值因子
                - rmw_freeshares: 盈利因子
                - cma_freeshares: 投资因子
    """
    if factor_type == 'FAMA_3YZ':
        select = copy.deepcopy(FAMA_3_FACTOR)
    elif factor_type == 'FAMA_5YZ':
        select = copy.deepcopy(FAMA_5_FACTOR)
    else:
        raise NotImplementedError('目前只有5因子和3因子')
    select.insert(0, 'trade_date')

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    where = []

    conn = db_conn.get_tk_factors_conn()

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=T_FAMA_FACTOR_RETURN, select='max(tradedate)',
                                       where=f"tradedate <= '{end_date}'",
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{T_FAMA_FACTOR_RETURN}表为空，请检查")
            return
        where.append(f"trade_date = '{latest_date}'")
    else:
        where.append(f"trade_date >= '{start_date}'")
        where.append(f"trade_date <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=T_FAMA_FACTOR_RETURN, select=select, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{factor_type}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df = df.rename(columns={'trade_date': 'end_date'})
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_barra_factor_return(factor_type, start_date=None, end_date=None, only_latest=False):
    """
    获得BARRA因子收益率数据

    :param str factor_type: BARRA_SYZ-风格因子；BARRA_IYZ-行业因子；BARRA_FYZ-全因子
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame，前十个为风格因子，后面的为行业因子
                end_date: 日期
                beta: 贝塔
                momentum: 动量
                mkt_size: 规模
                earnyild: 盈利
                resvol: 残差波动
                growth: 成长
                btop: 账面市值比
                leverage: 杠杆
                liquidty: 流动性
                sizenl: 非线性市值
                bank: 银行
                transportation: 交通运输
                leiservice: 休闲服务
                media: 传媒
                utilities: 公用事业
                agriforest: 农林牧渔
                chem: 化工
                health: 医药生物 
                commetrade: 商业贸易 
                aerodef: 国防军工 
                houseapp: 家用电器
                conmat: 建筑材料 
                builddeco: 建筑装饰 
                realestate: 房地产 
                nonfermetal: 有色金属
                machiequip: 机械设备 
                auto: 汽车 
                electronics: 电子 
                eleceqp: 电气设备
                textile: 纺织服装 
                conglomerates: 综合 
                computer: 计算机 
                lightindus: 轻工制造
                telecom: 通信 
                mining: 采掘 
                ironsteel: 钢铁 
                nonbankfinan: 非银金融
                foodbever: 食品饮料
                country: 国家
    """
    if factor_type == 'BARRA_SYZ':
        select = copy.deepcopy(BARRA_STYLE_FACTOR)
    elif factor_type == 'BARRA_IYZ':
        select = copy.deepcopy(BARRA_INDUSTRY_FACTOR)
    elif factor_type == 'BARRA_FYZ':
        select = copy.deepcopy(BARRA_FULL_FACTOR)
    else:
        raise NotImplementedError('目前只有风格因子、行业因子和全因子')
    select.insert(0, 'tradedate')

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    where = []

    conn = db_conn.get_tk_factors_conn()

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=T_BARRA_FACTOR_RETURN, select='max(tradedate)',
                                       where=f"tradedate <= '{end_date}'",
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{T_BARRA_FACTOR_RETURN}表为空，请检查")
            return
        where.append(f"tradedate = '{latest_date}'")
    else:
        where.append(f"tradedate >= '{start_date}'")
        where.append(f"tradedate <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=T_BARRA_FACTOR_RETURN, select=select, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{factor_type}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df = df.rename(columns={'tradedate': 'end_date'})
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_combined_factor_return(factor_list, start_date, end_date):
    """
    输入一个因子列表，可能包含bond、Fama、Barra因子，输出其收益率

    :param list factor_list: 因子名称的列表
    :param str start_date: 开始时间，格式2019-01-01
    :param str end_date: 结束时间，格式2020-01-01
    :return:
    """
    read_list = ['end_date'] + factor_list

    bond_factor_ret = get_bond_factor_return('BOND_7YZ', start_date, end_date)
    fama_factor_ret = get_fama_factor_return('FAMA_5YZ', start_date, end_date)
    barra_factor_ret = get_barra_factor_return('BARRA_FYZ', start_date, end_date)
    all_factor_ret = bond_factor_ret.merge(fama_factor_ret, how='inner', on='end_date')
    all_factor_ret = all_factor_ret.merge(barra_factor_ret, how='inner', on='end_date')

    res_ret = all_factor_ret[read_list]
    return res_ret


"""
之前的API函数，目前都替换成最新的API函数，以下函数已弃用
"""


# @common_wrappers.rename_date_col_to_tj('trade_date')
# @common_wrappers.automatic_date_formatter_wrapper(
#     {'trade_date': {'fmt': time_tool.PY_FMT_WITH_N_DASH, 'class': 'normal_date'}})
# def get_bond_factor_return(factor_name, date_related_filter=None):
#     """
#     获取债券因子收益率数据
#
#     :param str factor_name: BOND_4YZ: 债券四因子；BOND_7YZ: 债券七因子
#     :param date_related_filter: 日期过滤条件，格式 [('trade_date', '2020-01-01', '2020-03-01')]
#     :return:
#     """
#     conn = db_conn.get_tk_factors_conn()
#     if factor_name == 'BOND_4YZ':
#         select = copy.deepcopy(BOND_4_FACTOR)
#     elif factor_name == 'BOND_7YZ':
#         select = copy.deepcopy(BOND_7_FACTOR)
#     else:
#         raise NotImplementedError('目前只有4因子和7因子')
#     select.insert(0, 'trade_date')
#     where = date_related_filter
#     df = my_mysql.read_v2(sfrom=T_BOND_FACTOR_RETURN, select=select, where=where, conn=conn)
#     conn.close()
#     return df


# @common_wrappers.rename_date_col_to_tj('trade_date')
# @common_wrappers.automatic_date_formatter_wrapper(
#     {'trade_date': {'fmt': time_tool.PY_FMT_WITH_N_DASH, 'class': 'normal_date'}})
# def get_fama_factor_return(factor_name, date_related_filter=None):
#     """
#     获得FAMA因子收益率数据
#
#     :param str factor_name: FAMA_3YZ：3因子；FAMA_5YZ：5因子
#     :param date_related_filter: 日期过滤条件，格式 [('trade_date', '2020-01-01', '2020-03-01')]
#     :return:
#     """
#     if factor_name == 'FAMA_3YZ':
#         select = copy.deepcopy(FAMA_3_FACTOR)
#     elif factor_name == 'FAMA_5YZ':
#         select = copy.deepcopy(FAMA_5_FACTOR)
#     else:
#         raise NotImplementedError('目前只有5因子和3因子')
#     select.insert(0, 'trade_date')
#     where = date_related_filter
#     conn = db_conn.get_tk_factors_conn()
#     df = my_mysql.read_v2(sfrom=T_FAMA_FACTOR_RETURN, select=select, where=where, conn=conn)
#     conn.close()
#     return df


# @common_wrappers.rename_date_col_to_tj('tradedate')
# @common_wrappers.automatic_date_formatter_wrapper(
#     {'tradedate': {'fmt': time_tool.PY_FMT_WITH_N_DASH, 'class': 'normal_date'}})
# def get_barra_factor_return(factor_name, date_related_filter=None):
#     """
#     获得BARRA因子收益率数据
#
#     :param str factor_name: BARRA_SYZ：风格因子；BARRA_FYZ：全因子
#     :param date_related_filter: 日期过滤条件，格式 [('tradedate', '2020-01-01', '2020-03-01')]
#     :return:
#     """
#     conn = db_conn.get_tk_factors_conn()
#     where = date_related_filter
#
#     if factor_name == 'BARRA_SYZ':
#         select = copy.deepcopy(BARRA_STYLE_FACTOR)
#     elif factor_name == 'BARRA_FYZ':
#         df = my_mysql.read_v2(sfrom=T_BARRA_FACTOR_RETURN, conn=conn, limit=1)
#         cols_needed = list(df.columns.values)
#         cols_needed.remove('tradedate')
#         cols_needed.remove('tutimestamp')
#         cols_needed.remove('tutime')
#         cols_needed.remove('titime')
#         cols_needed.remove('tid')
#         cols_needed.remove('r2')
#         select = cols_needed
#     else:
#         raise NotImplementedError('目前只有 S因子 和 F因子')
#     select.insert(0, 'tradedate')
#
#     df = my_mysql.read_v2(
#         sfrom=T_BARRA_FACTOR_RETURN, select=select, where=where, conn=conn
#     )
#     conn.close()
#     return df


if __name__ == '__main__':
    # aaa = get_bond_factor_return('BOND_7YZ', '2020-01-01', '2020-03-01', False)
    # bbb = get_fama_factor_return('FAMA_5YZ', '2020-01-01', '2020-03-01', True)
    ccc = get_barra_factor_return('BARRA_FYZ', '2020-01-01', '2020-03-01', True)
